package mmpathgen.callgraphanalysis;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

import mmpathgen.util.JRELibFilter;
import soot.G;
import soot.Scene;
import soot.SootClass;
import soot.SootMethod;
import soot.jimple.toolkits.annotation.purity.SootMethodFilter;
import soot.jimple.toolkits.callgraph.CallGraph;
import soot.jimple.toolkits.callgraph.Edge;
import soot.toolkits.graph.Orderer;
import soot.toolkits.graph.PseudoTopologicalOrderer;

public class CallGraphVisitor {

	public static List<SootMethod> findMainMethods() {
		List<SootMethod> heads = new ArrayList<SootMethod>();
		for (SootClass cls : Scene.v().getApplicationClasses()) {
			for (SootMethod mtd : cls.getMethods()) {
				if (mtd.getName().equals("main")) {
					heads.add(mtd);
					G.v().out.println("  |- will treat " + cls.getName() + "."
							+ mtd.getName());
				}
			}
		}
		return heads;
	}

	// visit call graph nodes by width first traversal
	public static List<SootMethod> visitNodesInTopologicalOrder(UnfilteredCallGraph ucg) {

		// construct pseudo topological order on filtered methods
		Orderer<SootMethod> o = new PseudoTopologicalOrderer<SootMethod>();

		List<SootMethod> mtdList = new ArrayList<SootMethod>();
		// iterate the call graph in order
		for (SootMethod mtd : o.newList(ucg, false))
			mtdList.add(mtd);

		return mtdList;
	}
	
	// visit call graph nodes by width first traversal
	public static List<SootMethod> visitNodesInTopologicalOrder(FilteredCallGraph fcg) {

		// construct pseudo topological order on filtered methods
		Orderer<SootMethod> o = new PseudoTopologicalOrderer<SootMethod>();

		List<SootMethod> mtdList = new ArrayList<SootMethod>();
		// iterate the call graph in order
		for (SootMethod mtd : o.newList(fcg, false))
			mtdList.add(mtd);

		return mtdList;
	}

	// visit call graph nodes by width first traversal in reverse order
	public static List<SootMethod> visitNodesInReverseTopologicalOrder(UnfilteredCallGraph ucg) {

		// construct reverse pseudo topological order on filtered methods
		Orderer<SootMethod> o = new PseudoTopologicalOrderer<SootMethod>();

		List<SootMethod> mtdList = new ArrayList<SootMethod>();

		// iterate the call graph in reverse order
		for (SootMethod mtd : o.newList(ucg, true)) {
			mtdList.add(mtd);
		}

		return mtdList;
	}
	
	// visit call graph nodes by width first traversal in reverse order
	public static List<SootMethod> visitNodesInReverseTopologicalOrder(FilteredCallGraph fcg) {

		// construct reverse pseudo topological order on filtered methods
		Orderer<SootMethod> o = new PseudoTopologicalOrderer<SootMethod>();

		List<SootMethod> mtdList = new ArrayList<SootMethod>();

		// iterate the call graph in reverse order
		for (SootMethod mtd : o.newList(fcg, true)) {
			mtdList.add(mtd);
		}

		return mtdList;
	}

	// visit call graph edges by width first traversal
	public static List<Edge> visitEdgesInTopologicalOrder(CallGraph cg, boolean includeJRELib) {

		// Find main methods
		List<SootMethod> heads = findMainMethods();

		SootMethodFilter filter = new JRELibFilter();
		FilteredCallGraph fcg = new FilteredCallGraph(cg, filter, heads.iterator());

		// construct pseudo topological order on filtered methods
		Orderer<SootMethod> o = new PseudoTopologicalOrderer<SootMethod>();

		List<Edge> edgeList = new ArrayList<Edge>();
		Iterator<Edge> edgeIt = null;

		// iterate the call graph in order
		for (SootMethod srcMtd : o.newList(fcg, false)) {
			// System.out.println(srcMtd);
			edgeIt = cg.edgesOutOf(srcMtd);
			
			while (edgeIt.hasNext()) {
				Edge edge = edgeIt.next();
				if(includeJRELib)
					edgeList.add(edge);
				else if(!includeJRELib && fcg.getNodes().contains(edge.getTgt()))
					edgeList.add(edge);
				else // exclude JRE library and the target of the edge is an JRE library method 
					;
			}
		}

		return edgeList;
	}

	// visit call graph edges by width first traversal in reverse order
	public static List<Edge> visitEdgesInReverseTopologicalOrder(CallGraph cg,
			boolean includeJRELib) {

		// Find main methods
		List<SootMethod> heads = findMainMethods();

		SootMethodFilter filter = new JRELibFilter();
		FilteredCallGraph fcg = new FilteredCallGraph(cg, filter, heads.iterator());

		// construct pseudo topological order on filtered methods
		Orderer<SootMethod> o = new PseudoTopologicalOrderer<SootMethod>();

		List<Edge> edgeList = new ArrayList<Edge>();
		Iterator<Edge> edgeIt = null;

		// iterate the call graph in order
		for (SootMethod srcMtd : o.newList(fcg, true)) {
			// System.out.println(srcMtd);
			edgeIt = cg.edgesOutOf(srcMtd);
			
			while (edgeIt.hasNext()) {
				Edge edge = edgeIt.next();
				if(includeJRELib)
					edgeList.add(edge);
				else if(!includeJRELib && fcg.getNodes().contains(edge.getTgt()))
					edgeList.add(edge);
				else // exclude JRE library and the target of the edge is an JRE library method 
					;
			}
		}

		return edgeList;
	}

	public static List<SootMethod> getPossibleQuiescentMethods(CallGraph cg) {
		List<SootMethod> quiescentMtdList = new ArrayList<SootMethod>();

		// Find main methods
		List<SootMethod> heads = CallGraphVisitor.findMainMethods();

		SootMethodFilter filter = new JRELibFilter();
		FilteredCallGraph fcg = new FilteredCallGraph(cg, filter, heads.iterator());

		quiescentMtdList = fcg.getTails();

		return quiescentMtdList;
	}

}